#include "KeeperLogStore.h"

#include <cassert>

#include "nuraft.hxx"

namespace keeper_server {

KeeperLogStore::KeeperLogStore() : start_idx_(1) {
  NuRaftBufferPtr buf = NuRaftBuffer::alloc(sz_ulong);
  logs_[0] = nuraft::cs_new<NuRaftLogEntry>(0, buf);
}

KeeperLogStore::~KeeperLogStore() {}

NuRaftLogEntryPtr KeeperLogStore::make_clone(const NuRaftLogEntryPtr& entry) {
  NuRaftLogEntryPtr clone = nuraft::cs_new<NuRaftLogEntry>(
      entry->get_term(), NuRaftBuffer::clone(entry->get_buf()),
      entry->get_val_type());
  return clone;
}

ulong KeeperLogStore::next_slot() const {
  std::lock_guard<std::mutex> l(logs_lock_);
  // Exclude the dummy entry.
  return start_idx_ + logs_.size() - 1;
}

ulong KeeperLogStore::start_index() const { return start_idx_; }

NuRaftLogEntryPtr KeeperLogStore::last_entry() const {
  ulong next_idx = next_slot();
  std::lock_guard<std::mutex> l(logs_lock_);
  auto entry = logs_.find(next_idx - 1);
  if (entry == logs_.end()) {
    entry = logs_.find(0);
  }

  return make_clone(entry->second);
}

ulong KeeperLogStore::append(NuRaftLogEntryPtr& entry) {
  NuRaftLogEntryPtr clone = make_clone(entry);

  std::lock_guard<std::mutex> l(logs_lock_);
  size_t idx = start_idx_ + logs_.size() - 1;
  logs_[idx] = clone;
  return idx;
}

void KeeperLogStore::write_at(ulong index, NuRaftLogEntryPtr& entry) {
  NuRaftLogEntryPtr clone = make_clone(entry);

  // Discard all logs equal to or greater than `index.
  std::lock_guard<std::mutex> l(logs_lock_);
  auto itr = logs_.lower_bound(index);
  while (itr != logs_.end()) {
    itr = logs_.erase(itr);
  }
  logs_[index] = clone;
}

nuraft::ptr<NuRaftLogEntries> KeeperLogStore::log_entries(ulong start,
                                                          ulong end) {
  nuraft::ptr<NuRaftLogEntries> ret =
      nuraft::cs_new<std::vector<NuRaftLogEntryPtr>>();

  ret->resize(end - start);
  ulong cc = 0;
  for (ulong ii = start; ii < end; ++ii) {
    NuRaftLogEntryPtr src = nullptr;
    {
      std::lock_guard<std::mutex> l(logs_lock_);
      auto entry = logs_.find(ii);
      if (entry == logs_.end()) {
        entry = logs_.find(0);
        assert(0);
      }
      src = entry->second;
    }
    (*ret)[cc++] = make_clone(src);
  }
  return ret;
}

nuraft::ptr<NuRaftLogEntries> KeeperLogStore::log_entries_ext(
    ulong start, ulong end, nuraft::int64 batch_size_hint_in_bytes) {
  nuraft::ptr<NuRaftLogEntries> ret = nuraft::cs_new<NuRaftLogEntries>();

  if (batch_size_hint_in_bytes < 0) {
    return ret;
  }

  size_t accum_size = 0;
  for (ulong ii = start; ii < end; ++ii) {
    NuRaftLogEntryPtr src = nullptr;
    {
      std::lock_guard<std::mutex> l(logs_lock_);
      auto entry = logs_.find(ii);
      if (entry == logs_.end()) {
        entry = logs_.find(0);
        assert(0);
      }
      src = entry->second;
    }
    ret->push_back(make_clone(src));
    accum_size += src->get_buf().size();
    if (batch_size_hint_in_bytes &&
        accum_size >= (ulong)batch_size_hint_in_bytes)
      break;
  }
  return ret;
}

NuRaftLogEntryPtr KeeperLogStore::entry_at(ulong index) {
  NuRaftLogEntryPtr src = nullptr;
  {
    std::lock_guard<std::mutex> l(logs_lock_);
    auto entry = logs_.find(index);
    if (entry == logs_.end()) {
      entry = logs_.find(0);
    }
    src = entry->second;
  }
  return make_clone(src);
}

ulong KeeperLogStore::term_at(ulong index) {
  ulong term = 0;
  {
    std::lock_guard<std::mutex> l(logs_lock_);
    auto entry = logs_.find(index);
    if (entry == logs_.end()) {
      entry = logs_.find(0);
    }
    term = entry->second->get_term();
  }
  return term;
}

NuRaftBufferPtr KeeperLogStore::pack(ulong index, nuraft::int32 cnt) {
  std::vector<NuRaftBufferPtr> logs;

  size_t size_total = 0;
  for (ulong ii = index; ii < index + cnt; ++ii) {
    NuRaftLogEntryPtr le = nullptr;
    {
      std::lock_guard<std::mutex> l(logs_lock_);
      le = logs_[ii];
    }
    assert(le.get());
    NuRaftBufferPtr buf = le->serialize();
    size_total += buf->size();
    logs.push_back(buf);
  }

  NuRaftBufferPtr buf_out = nuraft::buffer::alloc(
      sizeof(nuraft::int32) + cnt * sizeof(nuraft::int32) + size_total);
  buf_out->pos(0);
  buf_out->put((nuraft::int32)cnt);

  for (auto& entry : logs) {
    NuRaftBufferPtr& bb = entry;
    buf_out->put((nuraft::int32)bb->size());
    buf_out->put(*bb);
  }
  return buf_out;
}

void KeeperLogStore::apply_pack(ulong index, NuRaftBuffer& pack) {
  pack.pos(0);
  nuraft::int32 num_logs = pack.get_int();

  for (nuraft::int32 ii = 0; ii < num_logs; ++ii) {
    ulong cur_idx = index + ii;
    nuraft::int32 buf_size = pack.get_int();

    NuRaftBufferPtr buf_local = NuRaftBuffer::alloc(buf_size);
    pack.get(buf_local);

    NuRaftLogEntryPtr le = NuRaftLogEntry::deserialize(*buf_local);
    {
      std::lock_guard<std::mutex> l(logs_lock_);
      logs_[cur_idx] = le;
    }
  }

  {
    std::lock_guard<std::mutex> l(logs_lock_);
    auto entry = logs_.upper_bound(0);
    if (entry != logs_.end()) {
      start_idx_ = entry->first;
    } else {
      start_idx_ = 1;
    }
  }
}

bool KeeperLogStore::compact(ulong last_log_index) {
  std::lock_guard<std::mutex> l(logs_lock_);
  for (ulong ii = start_idx_; ii <= last_log_index; ++ii) {
    auto entry = logs_.find(ii);
    if (entry != logs_.end()) {
      logs_.erase(entry);
    }
  }

  // WARNING:
  //   Even though nothing has been erased,
  //   we should set `start_idx_` to new index.
  if (start_idx_ <= last_log_index) {
    start_idx_ = last_log_index + 1;
  }
  return true;
}

void KeeperLogStore::close() { std::cout << "close" << std::endl; }

}  // namespace keeper_server
